from core.function import Function
from core.cuda import cuda
import utils

class Add(Function):
    def __init__(self):
        """
        初始化 加法 函数。
        """
        self.x1_shape = None
        self.x0_shape = None

    def forward(self, x0, x1):
        """
        加法 函数的前向传播。

        Args:
            x0 (Variable): 输入变量。
            x1 (Variable): 输入变量。

        Returns:
            Variable: 输出变量。
        """
        self.x0_shape, self.x1_shape = x0.shape, x1.shape
        x0 = cuda.to_array(x0)
        x1 = cuda.to_array(x1)
        y = x0 + x1
        return y

    def backward(self, gy):
        """
        加法 函数的反向传播。

        Args:
            gy (Variable): 相对于输出的梯度。

        Returns:
            tuple: 包含相对于输入 x0 和 x1 的梯度的元组。
        """
        gx0, gx1 = gy, gy
        # 加法运算的方向传播就是把上游的导数原封不动地“传走”
        # 广播
        # 如果在正向传播中进行了广播，就说明输入的 x0 和 x1 在形状上是不同的，
        # 此时应该进行广播的反向传播计算，为此需要求梯度 gx0 的和，使其变为 x0 的形状，同时还要处理 gx1
        if self.x0_shape != self.x1_shape:
            gx0 = utils.functions_collect.sum_to(gx0, self.x0_shape)
            gx1 = utils.functions_collect.sum_to(gx1, self.x1_shape)
        return gx0, gx1

    def add(self, x0, x1):
        """
        执行 加法 操作。

        Args:
            x0 (Variable): 输入变量。
            x1 (Variable): 输入变量。

        Returns:
            Variable: 输出变量。
        """
        return self(x0, x1)
